import pandas as pd
import os

from es_search import search_and_save_results

os.environ["CUDA_VISIBLE_DEVICES"] = "5"
import sys

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader
from langchain.embeddings.huggingface import HuggingFaceEmbeddings

file_path = '/home/zhengzhenzhuang/liujian/project/policy-model/data/name.txt'
embedding_path = '/home/zhengzhenzhuang/liujian/model/zpoint_large_embedding_zh'


def create_docsearch_index(file_path, embedding_path, query):
    # 加载文档  
    loader = TextLoader(file_path)
    documents = loader.load()

    # 分割文本  
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=0, chunk_overlap=0, separators=["-----"])
    texts = text_splitter.split_documents(documents)
    # print("texts",texts)
    # 创建嵌入  
    embeddings = HuggingFaceEmbeddings(model_name=embedding_path)

    # 构建搜索索引  
    docsearch = FAISS.from_documents(texts, embeddings)
    results = docsearch.similarity_search(query, k=10)

    return results[0].page_content, results

print("加载excel")
df = pd.read_excel("/home/zhengzhenzhuang/liujian/project/policy-model/search/result.xlsx")
data_list = df.values.tolist()

n = 0
for mylist in data_list:

    if "nan" == f"{mylist[3]}":
        n = n + 1
        question = mylist[1]
        if len(question)<5:
            continue
        search_and_save_results(question, None, None)

        print(mylist[1])
        a, b = create_docsearch_index(file_path, embedding_path, question)
        mylist[3]=a
        mylist[4]=b


        if n % 10 == 0:
            print(n)
            print("excel")
            df_from_list = pd.DataFrame(data_list)
            df_from_list.to_excel('/home/zhengzhenzhuang/liujian/project/policy-model/search/result-2.xlsx',
                                  index=False)
        print(n)
        print()

df_from_list = pd.DataFrame(data_list)
df_from_list.to_excel('/home/zhengzhenzhuang/liujian/project/policy-model/search/result-2.xlsx', index=False)
